#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys
import os
import cv2
from pathlib import Path
import numpy as np
import argparse
import pdb
import tsm_caffe_float as caffe

SUPPORT_IMGS = ['.jpg', '.jpeg', '.png', '.bmp'] 

def __build_arg_parser():
    parser = argparse.ArgumentParser(prog=f'{sys.argv[0]}'
            , description='To transform image to .npy files')

    parser.add_argument('files'
            , type = str
            , nargs = '+'
            , help = 'file list')
    return parser

def read_img(imgPath):
    f = str(imgPath.resolve())
    imgRaw = cv2.imread(f)
    # 这里是0-255之间的np.array
    # pdb.set_trace()
    return imgRaw

def __tran_to_npy(srcPath):
    f = str(srcPath.resolve())
    imgRaw = caffe.io.load_image(f)
    # pdb.set_trace()
   
    h = imgRaw.shape[0]
    w = imgRaw.shape[1]
    im_shrink = 640.0 / max(h, w)
    image = cv2.resize(imgRaw, None, None
            , fx=im_shrink, fy=im_shrink, interpolation=cv2.INTER_LINEAR)
    # pdb.set_trace()

    h = image.shape[0]
    w = image.shape[1]
    shape = (1, 3, h, w)
    transformer = caffe.io.Transformer({'data': shape})
    transformer.set_transpose('data', (2, 0, 1))
    transformer.set_mean('data', np.array([104, 117, 123])) # mean pixel
    transformer.set_raw_scale('data', 255)  # the reference model operates on images in [0,255] range instead of [0,1]
    transformer.set_channel_swap('data', (2, 1, 0))  # the reference model has channels in BGR order instead of RGB
    return transformer.preprocess('data', image)

def __check_file_path(imgPath):
    if not imgPath.exists():
        print (f'{imgPath} not found')
        sys.exit(1)

    if not imgPath.is_file():
        print (f'{imgPath} is not a file')
        sys.exit(1)

    if imgPath.suffix not in SUPPORT_IMGS:
        print (f'Only support ${SUPPORT_IMGS} files')
        sys.exit(1)

if __name__ == '__main__':
    parser = __build_arg_parser()
    args = parser.parse_args()

    for f in args.files:
        srcPath = Path(f)
        __check_file_path(srcPath)
        # read_img(srcPath)
        dstImg = __tran_to_npy(srcPath)
        dstPath = srcPath.with_suffix('.npy')
        print (f'translate {srcPath} to {dstPath}')
        np.save(str(dstPath.resolve()), dstImg)
